{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 将单机模型训练代码调整成 SecretFlow 联邦学习训练代码教程"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 引言\n",
    "#### 背景\n",
    "随着数据隐私问题日益受到重视，并且实际业务场景的需求不断演变，联邦学习作为一种特殊的深度学习形式开始迅速兴起。它以一种创新的方式解决了传统集中式训练的数据隐私问题，能够在保护用户隐私的前提下，有效地训练机器学习模型。在这个背景下，将大量的单机模型迁移到联邦学习框架下显得尤为重要和必要。\n",
    "\n",
    "将单机模型迁移到联邦学习框架下有着诸多优势。首先，它允许利用分散的数据资源，将数据存储在本地，避免了敏感信息集中存储的风险。这种分布式的学习方式不仅可以保护数据的隐私安全，还能够提高数据的安全性和可信度。其次，联邦学习将模型的训练过程分布在多个设备或者数据中心，因此能够增加模型的多样性和适应性。每个参与方在本地训练模型时，可以利用自身的数据特点和分布式计算资源，从而训练出更加贴合实际业务场景的模型。\n",
    "\n",
    "因此，将大量的单机模型迁移到联邦学习框架下不仅有助于保护数据隐私，还能够提高模型的训练效率和性能表现。这种迁移将成为未来深度学习领域的重要趋势，为各行业带来更多的创新和发展机会。\n",
    "#### 教程内容\n",
    "隐语联邦学习模块具有优异的兼容性，因此可以直接将基于 PyTorch 或 Tensorflow 的单机模型中的部分函数进行封装，并在 SecretFlow 中直接使用封装好的函数。那么具体是如何在水平联邦场景中将基于 PyTorch 或 Tensorflow 单机模型代码调整成 SecretFlow 的代码呢？\n",
    "\n",
    "除了前期数据准备工作外，本教程将基于 PyTorch 和基于 Tensorflow 分为两部分回答上述问题，并分别在每部分给出模型在单机框架和隐语框架下的实现，通过对比证明隐语框架下代码的可迁移性和易用性。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 前期数据准备\n",
    "\n",
    "下载数据集并解压"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import requests\n",
    "import tarfile\n",
    "import tempfile\n",
    "\n",
    "# Create a temporary folder\n",
    "_temp_dir = tempfile.mkdtemp()\n",
    "\n",
    "# Download file\n",
    "url = \"https://secretflow-data.oss-accelerate.aliyuncs.com/datasets/tf_flowers/flower_photos.tgz\"\n",
    "save_path = os.path.join(_temp_dir, \"flower_photos.tgz\")\n",
    "\n",
    "response = requests.get(url)\n",
    "with open(save_path, \"wb\") as f:\n",
    "    f.write(response.content)\n",
    "\n",
    "# Extract the file\n",
    "extract_folder = os.path.join(_temp_dir, \"flower_photos\")\n",
    "os.makedirs(extract_folder, exist_ok=True)\n",
    "\n",
    "with tarfile.open(save_path, \"r:gz\") as tar:\n",
    "    tar.extractall(path=extract_folder)\n",
    "\n",
    "path_to_flower_dataset = extract_folder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 基于PyTorch的迁移教程\n",
    "\n",
    "#### 模型在 PyTorch 下的实现\n",
    "首先我们给出单机模式下，基于 PyTorch 定义和训练神经网络模型的过程。对于数据，我们将其加载成 PyTorch \n",
    "中的Dataloader\n",
    "\n",
    "##### 构建 PyTorch 中 Dataloader 对象\n",
    "在单机模型中需要将数据集从文件夹加载成 torch 中的 Dataloader "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "import numpy as np\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data.sampler import SubsetRandomSampler\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "# parameter\n",
    "batch_size = 32\n",
    "shuffle = True\n",
    "random_seed = 1234\n",
    "train_split = 0.8\n",
    "\n",
    "# Define dataset\n",
    "flower_transform = transforms.Compose(\n",
    "    [\n",
    "        transforms.Resize((180, 180)),\n",
    "        transforms.ToTensor(),\n",
    "    ]\n",
    ")\n",
    "flower_dataset = datasets.ImageFolder(\n",
    "    path_to_flower_dataset, transform=flower_transform\n",
    ")\n",
    "dataset_size = len(flower_dataset)\n",
    "# Define sampler\n",
    "\n",
    "indices = list(range(dataset_size))\n",
    "if shuffle:\n",
    "    np.random.seed(random_seed)\n",
    "    np.random.shuffle(indices)\n",
    "split = int(np.floor(train_split * dataset_size))\n",
    "train_indices, val_indices = indices[:split], indices[split:]\n",
    "train_sampler = SubsetRandomSampler(train_indices)\n",
    "valid_sampler = SubsetRandomSampler(val_indices)\n",
    "\n",
    "# Define databuilder\n",
    "train_loader = DataLoader(flower_dataset, batch_size=batch_size, sampler=train_sampler)\n",
    "valid_loader = DataLoader(flower_dataset, batch_size=batch_size, sampler=valid_sampler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x.shape = torch.Size([32, 3, 180, 180])\n",
      "y.shape = torch.Size([32])\n"
     ]
    }
   ],
   "source": [
    "x, y = next(iter(train_loader))\n",
    "print(f\"x.shape = {x.shape}\")\n",
    "print(f\"y.shape = {y.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 基于 PyTorch 定义单机模型"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 定义模型架构"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "\n",
    "\n",
    "class ConvRGBNet_torch(nn.Module):\n",
    "    def __init__(self, *args, **kwargs) -> None:\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.network = nn.Sequential(\n",
    "            nn.Conv2d(\n",
    "                in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1\n",
    "            ),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(2, 2),\n",
    "            nn.Conv2d(16, 16, kernel_size=3, stride=1, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(2, 2),\n",
    "            nn.Flatten(),\n",
    "            nn.Linear(16 * 45 * 45, 128),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(128, 5),\n",
    "        )\n",
    "\n",
    "    def forward(self, xb):\n",
    "        return self.network(xb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 定义损失函数和优化器\n",
    "\n",
    "我们使用损失函数监督神经网络的训练，并且通过对 optimizer （优化器）指定超参数在训练过程中自动地调整模型参数，以使模型在给定的训练数据上达到最佳的性能表现。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize\n",
    "model_torch = ConvRGBNet_torch()\n",
    "\n",
    "# Define the loss function\n",
    "loss_model_torch = nn.CrossEntropyLoss()\n",
    "\n",
    "# Define the optimizer\n",
    "optimizer_model_torch = torch.optim.SGD(params=model_torch.parameters(), lr=0.01)\n",
    "optimizer_model_torch.zero_grad()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 训练模型\n",
    "\n",
    "一切就绪，我们训练模型。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, Iteration 30, Loss: 0.14022794975123057\n",
      "Epoch 2, Iteration 30, Loss: 0.0032508344719341647\n",
      "Epoch 3, Iteration 30, Loss: 0.0017656068026553839\n",
      "Epoch 4, Iteration 30, Loss: 0.0012027082809557518\n",
      "Epoch 5, Iteration 30, Loss: 0.0009076318300988835\n",
      "Finished Training\n"
     ]
    }
   ],
   "source": [
    "num_epochs = 5\n",
    "batch_num = len(train_loader)\n",
    "\n",
    "model_torch.train()  # prep model for training\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    train_loss = 0.0\n",
    "\n",
    "    for batch_idx, data in enumerate(train_loader):\n",
    "        inputs, labels = data\n",
    "        optimizer_model_torch.zero_grad()\n",
    "\n",
    "        outputs = model_torch(inputs)\n",
    "        loss = loss_model_torch(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer_model_torch.step()\n",
    "\n",
    "        train_loss += loss.item()\n",
    "        if (batch_idx + 1) % batch_num == 0:\n",
    "            print(\n",
    "                f\"Epoch {epoch+1}, Iteration {batch_idx+1}, Loss: {train_loss/batch_num}\"\n",
    "            )\n",
    "            # train_loss = 0.0\n",
    "\n",
    "print(\"Finished Training\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 测试模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy on test set: 1.0\n"
     ]
    }
   ],
   "source": [
    "# valid_loader is the data loader for the test dataset, and model_torch is the trained model\n",
    "model_torch.eval()\n",
    "total = 0\n",
    "correct = 0\n",
    "with torch.no_grad():\n",
    "    for images, labels in valid_loader:\n",
    "        outputs = model_torch(images)\n",
    "        _, predicted = torch.max(outputs.data, 1)\n",
    "        total += labels.size(0)\n",
    "        correct += (predicted == labels).sum().item()\n",
    "\n",
    "accuracy = correct / total\n",
    "print(f\"Accuracy on test set: {accuracy}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 基于 PyTorch 的隐语联邦学习模型迁移\n",
    "##### 概述 \n",
    "基于 PyTorch 从单机模型到联邦学习模型，主要包含以下步骤：\n",
    "- 添加联邦学习中的参与方\n",
    "- 修改数据集的处理逻辑\n",
    "- 修改模型的继承类\n",
    "- 根据需要决定是否对 metric 、 optimizer 和 loss fuction 进行包装，并使用包装后的函数\n",
    "\n",
    "接下来，我们将结合实际代码具体讲解这些步骤。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 环境设置\n",
    "添加联邦学习中的参与方，并对各个参与方进行初始化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The version of SecretFlow: 1.4.0.dev24011601\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-03-05 14:20:55,915\tINFO worker.py:1538 -- Started a local Ray instance.\n"
     ]
    }
   ],
   "source": [
    "import secretflow as sf\n",
    "\n",
    "# Check the version of your SecretFlow\n",
    "print('The version of SecretFlow: {}'.format(sf.__version__))\n",
    "\n",
    "# In case you have a running secretflow runtime already.\n",
    "sf.shutdown()\n",
    "sf.init(['alice', 'bob', 'charlie'], address=\"local\", log_to_driver=False)\n",
    "alice, bob, charlie = sf.PYU('alice'), sf.PYU('bob'), sf.PYU('charlie')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 封装单机模式下的数据处理逻辑\n",
    "同样地，在联邦学习中我们也需要对数据进行预处理，使之符合模型的输入，所以参考在[SecretFlow 中使用自定义 DataBuilder (Torch)构建 dataset builder](https://www.secretflow.org.cn/docs/secretflow/latest/zh-Hans/tutorial/CustomDataLoaderTorch)，我们选择文件夹路径作为参数，并且封装单机模式下的数据处理逻辑，最后返回 (data_set，steps_per_epoch)的结果，封装代码如下："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_dataset_builder(\n",
    "    batch_size=32,\n",
    "    train_split=0.8,\n",
    "    shuffle=True,\n",
    "    random_seed=1234,\n",
    "):\n",
    "    def dataset_builder(x, stage=\"train\"):\n",
    "        \"\"\" \"\"\"\n",
    "        import math\n",
    "\n",
    "        import numpy as np\n",
    "        from torch.utils.data import DataLoader\n",
    "        from torch.utils.data.sampler import SubsetRandomSampler\n",
    "        from torchvision import datasets, transforms\n",
    "\n",
    "        # Define dataset\n",
    "        flower_transform = transforms.Compose(\n",
    "            [\n",
    "                transforms.Resize((180, 180)),\n",
    "                transforms.ToTensor(),\n",
    "            ]\n",
    "        )\n",
    "        flower_dataset = datasets.ImageFolder(x, transform=flower_transform)\n",
    "        dataset_size = len(flower_dataset)\n",
    "        # Define sampler\n",
    "\n",
    "        indices = list(range(dataset_size))\n",
    "        if shuffle:\n",
    "            np.random.seed(random_seed)\n",
    "            np.random.shuffle(indices)\n",
    "        split = int(np.floor(train_split * dataset_size))\n",
    "        train_indices, val_indices = indices[:split], indices[split:]\n",
    "        train_sampler = SubsetRandomSampler(train_indices)\n",
    "        valid_sampler = SubsetRandomSampler(val_indices)\n",
    "\n",
    "        # Define databuilder\n",
    "        train_loader = DataLoader(\n",
    "            flower_dataset, batch_size=batch_size, sampler=train_sampler\n",
    "        )\n",
    "        valid_loader = DataLoader(\n",
    "            flower_dataset, batch_size=batch_size, sampler=valid_sampler\n",
    "        )\n",
    "\n",
    "        # Return\n",
    "        if stage == \"train\":\n",
    "            train_step_per_epoch = len(train_loader)\n",
    "\n",
    "            return train_loader, train_step_per_epoch\n",
    "        elif stage == \"eval\":\n",
    "            eval_step_per_epoch = len(valid_loader)\n",
    "            return valid_loader, eval_step_per_epoch\n",
    "\n",
    "    return dataset_builder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 构建 dataset_builder_dict\n",
    "\n",
    "我们通过 dataset_builder_dict 为各个参与方传入封装数据处理逻辑的 create_dataset_builder 函数的参数。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# prepare dataset dict\n",
    "data_builder_dict = {\n",
    "    alice: create_dataset_builder(\n",
    "        batch_size=32,\n",
    "        train_split=0.8,\n",
    "        shuffle=False,\n",
    "        random_seed=1234,\n",
    "    ),\n",
    "    bob: create_dataset_builder(\n",
    "        batch_size=32,\n",
    "        train_split=0.8,\n",
    "        shuffle=False,\n",
    "        random_seed=1234,\n",
    "    ),\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 在隐语的框架下定义基于 PyTorch 的模型架构\n",
    "参考 PyTorch 单机模式下的模型，我们在隐语的框架下定义同样结构的模型。\n",
    "\n",
    "我们只需修改继承类将 **torch.nn.Module** 改为 **secretflow.ml.nn.fl.utils.BaseModule**，就可以完成模型架构的定义。\n",
    "\n",
    "从迁移过程可以看出，将单机模型在隐语框架下进行定义所进行的代码改动非常小，整体迁移非常方便，充分展现了隐语框架的易用性。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-03-05 14:20:57.981319: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory\n",
      "2024-03-05 14:20:58.547649: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory\n",
      "2024-03-05 14:20:58.547720: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory\n",
      "2024-03-05 14:20:58.547727: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.\n"
     ]
    }
   ],
   "source": [
    "from secretflow.ml.nn.utils import BaseModule\n",
    "\n",
    "\n",
    "class ConvRGBNet(BaseModule):\n",
    "    def __init__(self, *args, **kwargs) -> None:\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.network = nn.Sequential(\n",
    "            nn.Conv2d(\n",
    "                in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1\n",
    "            ),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(2, 2),\n",
    "            nn.Conv2d(16, 16, kernel_size=3, stride=1, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(2, 2),\n",
    "            nn.Flatten(),\n",
    "            nn.Linear(16 * 45 * 45, 128),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(128, 5),\n",
    "        )\n",
    "\n",
    "    def forward(self, xb):\n",
    "        return self.network(xb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from secretflow.ml.nn import FLModel\n",
    "from secretflow.security.aggregation import SecureAggregator\n",
    "from torch import nn, optim\n",
    "from torchmetrics import Accuracy, Precision\n",
    "from secretflow.ml.nn.fl.utils import metric_wrapper, optim_wrapper\n",
    "from secretflow.ml.nn.utils import TorchModel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 将单机模型下的 optimizer 包装（wrap）\n",
    "- 对 optimizer 进行包装的原因\n",
    "\n",
    "如前所述，我们需要使用 optimizer 在训练过程中自动地调整模型参数，以使模型在给定的训练数据上达到最佳的性能表现。同样地，联邦学习也可以通过 optimizer 实现模型更好的性能。但由于隐语的设计机制同 PyTorch 有些差异，因此在优化器需要指定参数时，需要做一次封装来保证两者的一致性。\n",
    "\n",
    "通过`optim_wrapper`进行优化器（optimizer）的包装，并且通过追根溯源，我们可以看到\n",
    "```python\n",
    "from secretflow.ml.nn.fl.utils import optim_wrapper\n",
    "```\n",
    "可知此函数来自于`secretflow.ml.nn.fl.utils`,阅读[源代码](https://github.com/secretflow/secretflow/blob/main/secretflow/ml/nn/fl/utils.py#L94)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def optim_wrapper(func, *args, **kwargs):\n",
    "    def wrapped_func(params):\n",
    "        return func(params, *args, **kwargs)\n",
    "\n",
    "    return wrapped_func"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "可以看到函数实际上都是通过传入一个需要包装的函数名称，位置参数和关键字参数对函数完成包装，通过关键字参数确保了指定的参数赋值，然后返回包装好的函数。\n",
    "因此\n",
    "```python\n",
    "optim_fn = optim_wrapper(optim.Adam, lr=1e-2)\n",
    "```\n",
    "实际上相当于调用\n",
    "```python\n",
    "optim.Adam(lr=1e-2)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 将单机模型下的 metric 包装（wrap）\n",
    "- 对 metric 进行包装的原因\n",
    "\n",
    "metric 在机器学习和深度学习中用于衡量模型的性能和表现，它们是评估模型在训练、验证或测试数据上的效果的标准。同样地，我们也希望使用 metric 衡量联邦学习模型的性能和表现。但由于隐语的设计机制同 PyTorch 有些差异，因此在衡量指标需要指定参数时，需要做一次封装来保证两者的一致性。\n",
    "\n",
    "通过`metric_wrapper`进行衡量指标（metric）的包装，并且通过追根溯源，我们可以看到\n",
    "```python\n",
    "from secretflow.ml.nn.fl.utils import metric_wrapper\n",
    "```\n",
    "可知此函数来自于`secretflow.ml.nn.fl.utils`，阅读[源代码](https://github.com/secretflow/secretflow/blob/main/secretflow/ml/nn/fl/utils.py#L94)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def metric_wrapper(func, *args, **kwargs):\n",
    "    def wrapped_func():\n",
    "        return func(*args, **kwargs)\n",
    "\n",
    "    return wrapped_func"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "可以看到函数实际上是通过传入一个需要包装的函数名称，位置参数和关键字参数对函数完成包装，通过关键字参数确保了指定的参数赋值，然后返回包装好的函数；与优化器包装类似\n",
    "因此\n",
    "```python\n",
    "metric_wrapper(Accuracy, task=\"multiclass\", num_classes=10, average='micro')\n",
    "```\n",
    "实际上相当于调用\n",
    "```python\n",
    "Accuracy(task=\"multiclass\", num_classes=10, average='micro')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 将单机模型下的 loss function 包装（wrap）\n",
    "参考 `optim_wrapper` 和 `metric_wrapper` 的定义方式，自定义 `loss_function_wrapper`对损失函数（loss function）进行包装。\n",
    "\n",
    "如果模型使用损失函数的默认参数，则不需要使用包装"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss_function_wrapper(func, *args, **kwargs):\n",
    "    def wrapped_func():\n",
    "        return func(*args, **kwargs)\n",
    "\n",
    "    return wrapped_func"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "得益于隐语的封装，并且根据前述的包装原理可知，我们实际上需要通过包装完成一个参数具体化的函数，所以在这里 `nn.CrossEntropyLoss` 是需要包装的函数，并且其参数取值，就是需要传入的参数值。因为在其默认参数取值设置中，  `reduction='mean'` , 此处我们试着将其修改为 ` reduction='sum'`。包装损失函数只需要写成："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_wrapper = loss_function_wrapper(nn.CrossEntropyLoss, reduction='sum')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 对单机模型使用包装（wrap）后的优化器（optimizer）、衡量指标（metric）和损失函数（loss function）\n",
    "在隐语框架下，使用者可以根据需要选择是否对优化器（optimizer）、衡量指标（metric）和损失函数（loss function）进行包装，从而更好地训练联邦学习模型，充分体现了隐语框架使用的便捷性。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Create proxy actor <class 'secretflow.security.aggregation.secure_aggregator._Masker'> with party alice.\n",
      "INFO:root:Create proxy actor <class 'secretflow.security.aggregation.secure_aggregator._Masker'> with party bob.\n"
     ]
    }
   ],
   "source": [
    "device_list = [alice, bob]\n",
    "aggregator = SecureAggregator(charlie, [alice, bob])\n",
    "# prepare model\n",
    "num_classes = 5\n",
    "\n",
    "input_shape = (180, 180, 3)\n",
    "# torch model\n",
    "loss_wrapper = loss_function_wrapper(nn.CrossEntropyLoss, reduction='sum')\n",
    "optim_fn = optim_wrapper(optim.Adam, lr=1e-3)\n",
    "model_def = TorchModel(\n",
    "    model_fn=ConvRGBNet,\n",
    "    loss_fn=loss_wrapper,\n",
    "    optim_fn=optim_fn,\n",
    "    metrics=[\n",
    "        metric_wrapper(\n",
    "            Accuracy, task=\"multiclass\", num_classes=num_classes, average='micro'\n",
    "        ),\n",
    "        metric_wrapper(\n",
    "            Precision, task=\"multiclass\", num_classes=num_classes, average='micro'\n",
    "        ),\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 定义 FLModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Create proxy actor <class 'secretflow.ml.nn.fl.backend.torch.strategy.fed_avg_w.PYUFedAvgW'> with party alice.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Create proxy actor <class 'secretflow.ml.nn.fl.backend.torch.strategy.fed_avg_w.PYUFedAvgW'> with party bob.\n"
     ]
    }
   ],
   "source": [
    "fed_model = FLModel(\n",
    "    device_list=device_list,\n",
    "    model=model_def,\n",
    "    aggregator=aggregator,\n",
    "    backend=\"torch\", # backend support ['tensorflow', 'torch']\n",
    "    strategy=\"fed_avg_w\",\n",
    "    random_seed=1234,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 训练和验证模型\n",
    "传入参与方的数据集路径，进行模型的训练和验证"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:FL Train Params: {'x': {PYURuntime(alice): '/tmp/tmp35frzq4v/flower_photos', PYURuntime(bob): '/tmp/tmp35frzq4v/flower_photos'}, 'y': None, 'batch_size': 32, 'batch_sampling_rate': None, 'epochs': 5, 'verbose': 1, 'callbacks': None, 'validation_data': {PYURuntime(alice): '/tmp/tmp35frzq4v/flower_photos', PYURuntime(bob): '/tmp/tmp35frzq4v/flower_photos'}, 'shuffle': False, 'class_weight': None, 'sample_weight': None, 'validation_freq': 1, 'aggregate_freq': 2, 'label_decoder': None, 'max_batch_size': 20000, 'prefetch_buffer_size': None, 'sampler_method': 'batch', 'random_seed': 1234, 'dp_spent_step_freq': 1, 'audit_log_dir': None, 'dataset_builder': {PYURuntime(alice): <function create_dataset_builder.<locals>.dataset_builder at 0x7f68a526df70>, PYURuntime(bob): <function create_dataset_builder.<locals>.dataset_builder at 0x7f68282573a0>}, 'wait_steps': 100, 'self': <secretflow.ml.nn.fl.fl_model.FLModel object at 0x7f6818c0f7f0>}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train Processing: :   0%|          | 0/30 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train Processing: :  93%|█████████▎| 28/30 [00:50<00:03,  1.89s/it]/opt/anaconda3/envs/limingbo_secretflow/lib/python3.8/site-packages/secretflow/ml/nn/metrics.py:59: UserWarning: Please pay attention to local metrics, global only do naive aggregation.\n",
      "  warnings.warn(\n",
      "2024-03-05 14:22:00.417964: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory\n",
      "2024-03-05 14:22:00.418059: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcublas.so.11'; dlerror: libcublas.so.11: cannot open shared object file: No such file or directory\n",
      "2024-03-05 14:22:00.418107: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcublasLt.so.11'; dlerror: libcublasLt.so.11: cannot open shared object file: No such file or directory\n",
      "2024-03-05 14:22:00.418152: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcufft.so.10'; dlerror: libcufft.so.10: cannot open shared object file: No such file or directory\n",
      "2024-03-05 14:22:00.419139: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcusparse.so.11'; dlerror: libcusparse.so.11: cannot open shared object file: No such file or directory\n",
      "2024-03-05 14:22:00.419199: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1934] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.\n",
      "Skipping registering GPU devices...\n",
      "Train Processing: :  93%|█████████▎| 28/30 [00:53<00:03,  1.91s/it, {'multiclassaccuracy': 0.96666664, 'multiclassprecision': 0.96666664, 'val_multiclassaccuracy': 1.0, 'val_multiclassprecision': 1.0}]\n",
      "Train Processing: :   0%|          | 0/30 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train Processing: :  93%|█████████▎| 28/30 [00:47<00:03,  1.69s/it, {'multiclassaccuracy': 1.0, 'multiclassprecision': 1.0, 'val_multiclassaccuracy': 1.0, 'val_multiclassprecision': 1.0}]\n",
      "Train Processing: :   0%|          | 0/30 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 3/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train Processing: :  93%|█████████▎| 28/30 [00:45<00:03,  1.62s/it, {'multiclassaccuracy': 1.0, 'multiclassprecision': 1.0, 'val_multiclassaccuracy': 1.0, 'val_multiclassprecision': 1.0}]\n",
      "Train Processing: :   0%|          | 0/30 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 4/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train Processing: :  93%|█████████▎| 28/30 [00:45<00:03,  1.61s/it, {'multiclassaccuracy': 1.0, 'multiclassprecision': 1.0, 'val_multiclassaccuracy': 1.0, 'val_multiclassprecision': 1.0}]\n",
      "Train Processing: :   0%|          | 0/30 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 5/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train Processing: :  93%|█████████▎| 28/30 [00:45<00:03,  1.61s/it, {'multiclassaccuracy': 1.0, 'multiclassprecision': 1.0, 'val_multiclassaccuracy': 1.0, 'val_multiclassprecision': 1.0}]\n"
     ]
    }
   ],
   "source": [
    "data = {\n",
    "    alice: path_to_flower_dataset,\n",
    "    bob: path_to_flower_dataset,\n",
    "}\n",
    "history = fed_model.fit(\n",
    "    data,\n",
    "    None,\n",
    "    validation_data=data,\n",
    "    epochs=5,\n",
    "    batch_size=32,\n",
    "    aggregate_freq=2,\n",
    "    sampler_method=\"batch\",\n",
    "    random_seed=1234,\n",
    "    dp_spent_step_freq=1,\n",
    "    dataset_builder=data_builder_dict,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "可以看到，联邦学习的训练代码顺利运行"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 小结\n",
    "基于 PyTorch 从单机模型到联邦学习模型，主要需要添加或修改以下几部分\n",
    "- 添加联邦学习中的参与方\n",
    "- 修改数据集的处理逻辑\n",
    "- 修改模型的继承类\n",
    "- 根据需要决定是否对 metric 、 optimizer 和 loss fuction 进行包装，并使用包装后的函数\n",
    "\n",
    "得益于隐语的封装，使用者不需要自己完成模型定义等代码的编写，只需要调用 Secretflow 中的函数，即可便捷完成模型的定义和使用等操作。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 基于 TensorFlow 的迁移教程\n",
    "#### 模型在 TensorFlow 下的实现\n",
    "首先我们给出单机模式下，基于 TensorFlow 定义和训练神经网络模型的过程。对于数据，我们将其加载成 TensorFlow 的 dataset 对象\n",
    "##### 构建 TensorFlow 中  tf.dataset  对象\n",
    "在单机模型中需要将数据集从文件夹加载成 TensorFlow 的 dataset 对象"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 1201 files belonging to 1 classes.\n",
      "Using 961 files for training.\n",
      "Using 240 files for validation.\n"
     ]
    }
   ],
   "source": [
    "import math\n",
    "import tensorflow as tf\n",
    "\n",
    "img_height = 180\n",
    "img_width = 180\n",
    "batch_size = 32\n",
    "# In this example, we use the TensorFlow interface for development.\n",
    "data_set = tf.keras.utils.image_dataset_from_directory(\n",
    "    path_to_flower_dataset,\n",
    "    validation_split=0.2,\n",
    "    subset=\"both\",\n",
    "    seed=123,\n",
    "    image_size=(img_height, img_width),\n",
    "    batch_size=batch_size,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_set = data_set[0]\n",
    "test_set = data_set[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'tensorflow.python.data.ops.dataset_ops.BatchDataset'> <class 'tensorflow.python.data.ops.dataset_ops.BatchDataset'>\n"
     ]
    }
   ],
   "source": [
    "print(type(train_set), type(test_set))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x.shape = (32, 180, 180, 3)\n",
      "y.shape = (32,)\n"
     ]
    }
   ],
   "source": [
    "x, y = next(iter(train_set))\n",
    "print(f\"x.shape = {x.shape}\")\n",
    "print(f\"y.shape = {y.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 基于 Tensorflow 定义单机模型\n",
    "- 定义模型的架构\n",
    "\n",
    "我们进行模型架构的定义，并且使用损失函数监督神经网络的训练，通过 optimizer （优化器）在训练过程中自动地调整模型参数，以使模型在给定的训练数据上达到最佳的性能表现。在 TensorFlow 里面，优化器（optimizer）、衡量指标（metric）和损失函数（loss function）可以通过在 compile 函数中通过参数指定值的方式实现。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow import keras\n",
    "\n",
    "# Create model\n",
    "num_classes = 5\n",
    "input_shape = (180, 180, 3)\n",
    "total_epochs = 10\n",
    "\n",
    "model_tensorflow = keras.Sequential(\n",
    "    [\n",
    "        keras.Input(shape=input_shape),\n",
    "        tf.keras.layers.Rescaling(1.0 / 255),\n",
    "        tf.keras.layers.Conv2D(32, 3, activation='relu'),\n",
    "        tf.keras.layers.MaxPooling2D(),\n",
    "        tf.keras.layers.Conv2D(32, 3, activation='relu'),\n",
    "        tf.keras.layers.MaxPooling2D(),\n",
    "        tf.keras.layers.Conv2D(32, 3, activation='relu'),\n",
    "        tf.keras.layers.MaxPooling2D(),\n",
    "        tf.keras.layers.Flatten(),\n",
    "        tf.keras.layers.Dense(128, activation='relu'),\n",
    "        tf.keras.layers.Dense(num_classes),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Compile model\n",
    "model_tensorflow.compile(\n",
    "    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
    "    optimizer='adam',\n",
    "    metrics=[\"accuracy\"],\n",
    ")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 训练和验证模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-------Start training-------\n",
      "Epoch 1/10\n",
      "31/31 [==============================] - 5s 132ms/step - loss: 0.0533 - accuracy: 0.9740 - val_loss: 0.0000e+00 - val_accuracy: 1.0000\n",
      "Epoch 2/10\n",
      "31/31 [==============================] - 4s 124ms/step - loss: 0.0000e+00 - accuracy: 1.0000 - val_loss: 0.0000e+00 - val_accuracy: 1.0000\n",
      "Epoch 3/10\n",
      "31/31 [==============================] - 4s 134ms/step - loss: 0.0000e+00 - accuracy: 1.0000 - val_loss: 0.0000e+00 - val_accuracy: 1.0000\n",
      "Epoch 4/10\n",
      "31/31 [==============================] - 4s 132ms/step - loss: 0.0000e+00 - accuracy: 1.0000 - val_loss: 0.0000e+00 - val_accuracy: 1.0000\n",
      "Epoch 5/10\n",
      "31/31 [==============================] - 4s 125ms/step - loss: 0.0000e+00 - accuracy: 1.0000 - val_loss: 0.0000e+00 - val_accuracy: 1.0000\n",
      "Epoch 6/10\n",
      "31/31 [==============================] - 4s 126ms/step - loss: 0.0000e+00 - accuracy: 1.0000 - val_loss: 0.0000e+00 - val_accuracy: 1.0000\n",
      "Epoch 7/10\n",
      "31/31 [==============================] - 4s 121ms/step - loss: 0.0000e+00 - accuracy: 1.0000 - val_loss: 0.0000e+00 - val_accuracy: 1.0000\n",
      "Epoch 8/10\n",
      "31/31 [==============================] - 4s 118ms/step - loss: 0.0000e+00 - accuracy: 1.0000 - val_loss: 0.0000e+00 - val_accuracy: 1.0000\n",
      "Epoch 9/10\n",
      "31/31 [==============================] - 4s 119ms/step - loss: 0.0000e+00 - accuracy: 1.0000 - val_loss: 0.0000e+00 - val_accuracy: 1.0000\n",
      "Epoch 10/10\n",
      "31/31 [==============================] - 4s 120ms/step - loss: 0.0000e+00 - accuracy: 1.0000 - val_loss: 0.0000e+00 - val_accuracy: 1.0000\n"
     ]
    }
   ],
   "source": [
    "# Model training and validation\n",
    "print('-------Start training-------')\n",
    "history = model_tensorflow.fit(\n",
    "    train_set,\n",
    "    validation_data=test_set,\n",
    "    batch_size=batch_size,\n",
    "    epochs=total_epochs,\n",
    "    verbose=True,\n",
    "    shuffle=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 基于 TensorFlow 的隐语联邦学习模型迁移\n",
    "##### 概述 \n",
    "基于 TensorFlow 从单机模型到联邦学习模型，主要包含以下步骤：\n",
    "- 添加联邦学习中的参与方\n",
    "- 修改数据集的处理逻辑\n",
    "- 进行模型的封装\n",
    "\n",
    "接下来，我们将结合实际代码具体讲解这些步骤。\n",
    "##### 环境设置\n",
    "添加联邦学习中的参与方，并初始化各个参与方"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The version of SecretFlow: 1.4.0.dev24011601\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-03-05 14:25:51,394\tINFO worker.py:1538 -- Started a local Ray instance.\n"
     ]
    }
   ],
   "source": [
    "import secretflow as sf\n",
    "\n",
    "# Check the version of your SecretFlow\n",
    "print('The version of SecretFlow: {}'.format(sf.__version__))\n",
    "\n",
    "# In case you have a running secretflow runtime already.\n",
    "sf.shutdown()\n",
    "sf.init(['alice', 'bob', 'charlie'], address=\"local\", log_to_driver=False)\n",
    "alice, bob, charlie = sf.PYU('alice'), sf.PYU('bob'), sf.PYU('charlie')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 封装单机模式下的数据处理逻辑\n",
    "\n",
    "同样地，在联邦学习中我们也需要对数据进行预处理，使之符合模型的输入，参考在[SecretFlow 中使用自定义 DataBuilder (TensorFlow)构建 dataset builder](https://www.secretflow.org.cn/docs/secretflow/latest/zh-Hans/tutorial/CustomDataLoaderTF)，我们选择文件夹路径作为参数，并且封装单机模式下的数据处理逻辑，最后返回 (data_set，steps_per_epoch)的结果"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_dataset_builder(\n",
    "    batch_size=32,\n",
    "):\n",
    "    def dataset_builder(folder_path, stage=\"train\"):\n",
    "        import math\n",
    "\n",
    "        import tensorflow as tf\n",
    "\n",
    "        img_height = 180\n",
    "        img_width = 180\n",
    "        data_set = tf.keras.utils.image_dataset_from_directory(\n",
    "            folder_path,\n",
    "            validation_split=0.2,\n",
    "            subset=\"both\",\n",
    "            seed=123,\n",
    "            image_size=(img_height, img_width),\n",
    "            batch_size=batch_size,\n",
    "        )\n",
    "        if stage == \"train\":\n",
    "            train_dataset = data_set[0]\n",
    "            train_step_per_epoch = math.ceil(len(data_set[0].file_paths) / batch_size)\n",
    "            return train_dataset, train_step_per_epoch\n",
    "        elif stage == \"eval\":\n",
    "            eval_dataset = data_set[1]\n",
    "            eval_step_per_epoch = math.ceil(len(data_set[1].file_paths) / batch_size)\n",
    "            return eval_dataset, eval_step_per_epoch\n",
    "\n",
    "    return dataset_builder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 构建 dataset_builder_dict\n",
    "我们通过 dataset_builder_dict 为各个参与方传入封装数据处理逻辑的 create_dataset_builder 函数的参数。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_builder_dict = {\n",
    "    alice: create_dataset_builder(\n",
    "        batch_size=32,\n",
    "    ),\n",
    "    bob: create_dataset_builder(\n",
    "        batch_size=32,\n",
    "    ),\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 在隐语的框架下定义基于 TensorFlow 的模型架构\n",
    "参考 TensorFlow 单机模式下的模型，我们在隐语的框架下定义同样结构的模型。\n",
    "\n",
    "如前所述，我们需要使用 optimizer 在训练过程中自动地调整模型参数，以使模型在给定的训练数据上达到最佳的性能表现；使用 metric 在机器学习和深度学习中用于衡量模型的性能和表现；使用损失函数监督神经网络的训练。同样地，联邦学习也可以通过使用优化器（optimizer）、衡量指标（metric）和损失函数（loss function） 实现模型更好的性能。在隐语框架下，优化器（optimizer）、衡量指标（metric）和损失函数（loss function）的使用方法和 TensorFlow 中的使用方法一致，也同样通过模型的 compile 函数实现。\n",
    "\n",
    "从迁移过程可以看出，代码修改幅度非常小，整体迁移过程非常方便，充分展现了隐语框架的易用性和便捷性。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_conv_flower_model(input_shape, num_classes, name='model'):\n",
    "    def create_model():\n",
    "        from tensorflow import keras\n",
    "\n",
    "        # Create model\n",
    "\n",
    "        model = keras.Sequential(\n",
    "            [\n",
    "                keras.Input(shape=input_shape),\n",
    "                tf.keras.layers.Rescaling(1.0 / 255),\n",
    "                tf.keras.layers.Conv2D(32, 3, activation='relu'),\n",
    "                tf.keras.layers.MaxPooling2D(),\n",
    "                tf.keras.layers.Conv2D(32, 3, activation='relu'),\n",
    "                tf.keras.layers.MaxPooling2D(),\n",
    "                tf.keras.layers.Conv2D(32, 3, activation='relu'),\n",
    "                tf.keras.layers.MaxPooling2D(),\n",
    "                tf.keras.layers.Flatten(),\n",
    "                tf.keras.layers.Dense(128, activation='relu'),\n",
    "                tf.keras.layers.Dense(num_classes),\n",
    "            ]\n",
    "        )\n",
    "        # Compile model\n",
    "        model.compile(\n",
    "            loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
    "            optimizer='adam',\n",
    "            metrics=[\"accuracy\"],\n",
    "        )\n",
    "        return model\n",
    "\n",
    "    return create_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "from secretflow.ml.nn import FLModel\n",
    "from secretflow.security.aggregation import SecureAggregator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Create proxy actor <class 'secretflow.security.aggregation.secure_aggregator._Masker'> with party alice.\n",
      "INFO:root:Create proxy actor <class 'secretflow.security.aggregation.secure_aggregator._Masker'> with party bob.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Create proxy actor <class 'secretflow.ml.nn.fl.backend.tensorflow.strategy.fed_avg_w.PYUFedAvgW'> with party alice.\n",
      "INFO:root:Create proxy actor <class 'secretflow.ml.nn.fl.backend.tensorflow.strategy.fed_avg_w.PYUFedAvgW'> with party bob.\n"
     ]
    }
   ],
   "source": [
    "device_list = [alice, bob]\n",
    "aggregator = SecureAggregator(charlie, [alice, bob])\n",
    "\n",
    "# prepare model\n",
    "num_classes = 5\n",
    "input_shape = (180, 180, 3)\n",
    "\n",
    "# keras model\n",
    "model = create_conv_flower_model(input_shape, num_classes)\n",
    "\n",
    "\n",
    "fed_model = FLModel(\n",
    "    device_list=device_list,\n",
    "    model=model,\n",
    "    aggregator=aggregator,\n",
    "    backend=\"tensorflow\",\n",
    "    strategy=\"fed_avg_w\",\n",
    "    random_seed=1234,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 训练和验证模型\n",
    "传入参与方的数据集路径，进行模型的训练和验证"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:FL Train Params: {'x': {PYURuntime(alice): '/tmp/tmp35frzq4v/flower_photos', PYURuntime(bob): '/tmp/tmp35frzq4v/flower_photos'}, 'y': None, 'batch_size': 32, 'batch_sampling_rate': None, 'epochs': 5, 'verbose': 1, 'callbacks': None, 'validation_data': {PYURuntime(alice): '/tmp/tmp35frzq4v/flower_photos', PYURuntime(bob): '/tmp/tmp35frzq4v/flower_photos'}, 'shuffle': False, 'class_weight': None, 'sample_weight': None, 'validation_freq': 1, 'aggregate_freq': 2, 'label_decoder': None, 'max_batch_size': 20000, 'prefetch_buffer_size': None, 'sampler_method': 'batch', 'random_seed': 1234, 'dp_spent_step_freq': 1, 'audit_log_dir': None, 'dataset_builder': {PYURuntime(alice): <function create_dataset_builder.<locals>.dataset_builder at 0x7f6818bd2160>, PYURuntime(bob): <function create_dataset_builder.<locals>.dataset_builder at 0x7f682824ae50>}, 'wait_steps': 100, 'self': <secretflow.ml.nn.fl.fl_model.FLModel object at 0x7f5be84413a0>}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train Processing: :   0%|          | 0/31 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train Processing: :  97%|█████████▋| 30/31 [00:14<00:00,  2.09it/s, {'loss': 0.06979854, 'accuracy': 0.9667013, 'val_loss': 0.0, 'val_accuracy': 1.0}]\n",
      "Train Processing: :   0%|          | 0/31 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train Processing: :  97%|█████████▋| 30/31 [00:13<00:00,  2.30it/s, {'loss': 0.0, 'accuracy': 1.0, 'val_loss': 0.0, 'val_accuracy': 1.0}]\n",
      "Train Processing: :   0%|          | 0/31 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 3/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train Processing: :  97%|█████████▋| 30/31 [00:13<00:00,  2.29it/s, {'loss': 0.0, 'accuracy': 1.0, 'val_loss': 0.0, 'val_accuracy': 1.0}]\n",
      "Train Processing: :   0%|          | 0/31 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 4/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train Processing: :  97%|█████████▋| 30/31 [00:13<00:00,  2.24it/s, {'loss': 0.0, 'accuracy': 1.0, 'val_loss': 0.0, 'val_accuracy': 1.0}]\n",
      "Train Processing: :   0%|          | 0/31 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 5/5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train Processing: :  97%|█████████▋| 30/31 [00:13<00:00,  2.25it/s, {'loss': 0.0, 'accuracy': 1.0, 'val_loss': 0.0, 'val_accuracy': 1.0}]\n"
     ]
    }
   ],
   "source": [
    "data = {\n",
    "    alice: path_to_flower_dataset,\n",
    "    bob: path_to_flower_dataset,\n",
    "}\n",
    "history = fed_model.fit(\n",
    "    data,\n",
    "    None,\n",
    "    validation_data=data,\n",
    "    epochs=5,\n",
    "    batch_size=32,\n",
    "    aggregate_freq=2,\n",
    "    sampler_method=\"batch\",\n",
    "    random_seed=1234,\n",
    "    dp_spent_step_freq=1,\n",
    "    dataset_builder=data_builder_dict,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 小结\n",
    "基于 TensorFlow 从单机模型到联邦学习模型，主要需要添加或修改以下几部分\n",
    "- 添加联邦学习中的参与方\n",
    "- 修改数据集的处理逻辑\n",
    "- 进行模型的封装\n",
    "\n",
    "得益于隐语的封装，使用者不需要自己进行大量的代码编写，只需要调用 Secretflow 中的函数，即可便捷完成模型的定义和使用等操作。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 总结\n",
    "本教程说明了使用者能够在 SecretFlow 隐语的框架下，体会到和单机模式下，使用 PyTorch 或 Tensorflow 编程几乎一致的联邦学习模型使用体验。充分展现了隐语框架具有易用性和便捷性等优点。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
